library(bartMachine)
library(grf)
library(caret)
x = read.csv('Xt.csv',header=FALSE)

testData<-function(n, binind) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  
  sigma_y = snr*sd(mu(x) + p(x)*alpha(x))
  
  xind=sample(1:nrow(x),n,replace=TRUE)
  n1=round(n/2)
  n0=n-n1
  zind=sample(1:n,size=n1)
  z=numeric(n)
  z[zind]=1
  het = rep(1,16)
  y = mu(x[xind,]) + z*alpha(x[xind,]) + het[as.numeric(x[xind,21])]*sigma_y*rnorm(n)
  dataframe = cbind(x[xind,c(1,3,10,14,15,21,24,43)], z = z, y = y)
  return(dataframe)
}
optimalITR <- function(binind) {
  
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  return(as.numeric(y1>y0))
}

truePAV<-function(That, binind) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  probs=sum(That)/length(That)
  PAV=mean(y1*That+y0*(1-That))
  return(PAV)
}

truePAPE<-function(That, binind) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  probs=sum(That)/length(That)
  PAPE=mean(y1*That+y0*(1-That))-probs*mean(y1)-(1-probs)*mean(y0)
  return(PAPE)
}

truePAPEp<-function(That, binind,plim) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  PAPE=mean(y1*That+y0*(1-That))-plim*mean(y1)-(1-plim)*mean(y0)
  return(PAPE)
}

truePAPD<-function(Thatf, Thatg, binind, plim) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  PAPEf=mean(y1*Thatf+y0*(1-Thatf))-plim*mean(y1)-(1-plim)*mean(y0)
  PAPEg=mean(y1*Thatg+y0*(1-Thatg))-plim*mean(y1)-(1-plim)*mean(y0)
  PAPD=PAPEf-PAPEg
  return(PAPD)
}

trueAUPEC<-function(tau, binind) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  
  
  y1 = mu(x) + alpha(x)
  y0 = mu(x)
  AUCT=0
  pprior=mean(y0)
  for (i in 1:length(tau)) {
    cutofftemp=quantile(tau,1-i/length(tau))
    Thattemp=as.numeric(tau>max(cutofftemp,0))
    AUCT=AUCT + (mean(y1*Thattemp+y0*(1-Thattemp))+pprior)*0.5/length(tau)
    pprior=mean(y1*Thattemp+y0*(1-Thattemp))
    }
  AUPEC=AUCT-1/2*mean(y1)-1/2*mean(y0)
  return(AUPEC)
}


testPAPEa<-function(n, binind, model, sep=0.5) {
  magnitude = binind
  if (magnitude ==0){
    effect_size = 1/3
  }
  
  if (magnitude == 1){
    effect_size = 2}
  
  snr=0.25
  
  beta0=-1
  beta1=3
  
  p = function(x){
    
    result = x[,1] + x[,43] + 0.3*(as.numeric(x[,10])-1)
    result = 1/(1+exp(beta0 + beta1*result))
    
  }
  
  mu = function(x){
    
    #result = 3*(x[,1] > 0 & x[,43]> 0) + (x[,1] < 0 & x[,43]> 0) + 0.5*(x[,1] > 0 & x[,43] < 0) + (as.numeric(x[,4])-3/2)*atan(x[,5]) + as.numeric(x[,3]) + 3*atan(qnorm(p(x))/3)
    result =  -sin(qnorm(p(x))) + x[,43]
    return(result)
    
  }
  
  
  
  alpha = function(x){
    
    result = as.numeric(x[,3] == 'zero')*as.numeric(x[,24] == 'black') + as.numeric(x[,14])-1 - (as.numeric(x[,15])-1)
    result = effect_size*result
    return(result)
  }
  sigma_y = snr*sd(mu(x) + p(x)*alpha(x))
  xind=sample(1:nrow(x),n,replace=TRUE)
  x_test=x[xind,c(1,3,10,14,15,21,24,43)]
  x_test0 = cbind(x_test, z=0)
  x_test1 = cbind(x_test, z=1)
  yp0 = predict(model, x_test0)
  yp1 = predict(model, x_test1)
  That = as.numeric(yp1>yp0)
  nf=round(n/2)
  nr=n-nf
  nr1 = round(sep*nr)
  nr0 = nr - nr1
  zind=sample(1:n,size=nf)
  z=numeric(n)
  z[zind]=That[zind]
  randsamps=setdiff(1:n,zind)
  randind1=sample(randsamps,size = nr1)
  randind0=setdiff(randsamps, randind1)
  z[randind1]=1
  y = mu(x[xind,]) + z*alpha(x[xind,]) + sigma_y*rnorm(n)
  probs=sum(That)/n
  PAPE = n/(n-1)*(1/nf*sum(z[zind]*y[zind]+(1-z[zind])*y[zind])-probs/nr1*sum(z[randsamps]*y[randsamps])-(1-probs)/nr0*sum((1-z[randsamps])*y[randsamps]))
  Sf=var(y[zind])
  S1=var(y[randind1])
  S0=var(y[randind0])
  SATE=1/nr1*sum(z[randsamps]*y[randsamps])-1/nr0*(sum((1-z[randsamps])*y[randsamps]))
  covarterm=1/n^2*(PAPE^2+2*(n-1)*PAPE*SATE*(2*probs-1)-(1-probs)*probs*n*SATE^2)
  varexp=n^2/(n-1)^2*(Sf/nf+S1*probs^2/nr1+S0*(1-probs)^2/nr0+covarterm)
  return(list(pape=PAPE,sd=sqrt(varexp)))
}

